﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace test0320
{
    public class AES
    {
        private readonly int _nk;// 密钥长度
        private readonly int _rounds;// 加密轮数
        private uint[] _roundKeys;//轮密钥

        /// <summary>
        /// s盒
        /// </summary>
        private static readonly byte[,] Sbox = new byte[16, 16]
        {
        {0x63,0x7C,0x77,0x7B,0xF2,0x6B,0x6F,0xC5,0x30,0x01,0x67,0x2B,0xFE,0xD7,0xAB,0x76},
        {0xCA,0x82,0xC9,0x7D,0xFA,0x59,0x47,0xF0,0xAD,0xD4,0xA2,0xAF,0x9C,0xA4,0x72,0xC0},
        {0xB7,0xFD,0x93,0x26,0x36,0x3F,0xF7,0xCC,0x34,0xA5,0xE5,0xF1,0x71,0xD8,0x31,0x15},
        {0x04,0xC7,0x23,0xC3,0x18,0x96,0x05,0x9A,0x07,0x12,0x80,0xE2,0xEB,0x27,0xB2,0x75},
        {0x09,0x83,0x2C,0x1A,0x1B,0x6E,0x5A,0xA0,0x52,0x3B,0xD6,0xB3,0x29,0xE3,0x2F,0x84},
        {0x53,0xD1,0x00,0xED,0x20,0xFC,0xB1,0x5B,0x6A,0xCB,0xBE,0x39,0x4A,0x4C,0x58,0xCF},
        {0xD0,0xEF,0xAA,0xFB,0x43,0x4D,0x33,0x85,0x45,0xF9,0x02,0x7F,0x50,0x3C,0x9F,0xA8},
        {0x51,0xA3,0x40,0x8F,0x92,0x9D,0x38,0xF5,0xBC,0xB6,0xDA,0x21,0x10,0xFF,0xF3,0xD2},
        {0xCD,0x0C,0x13,0xEC,0x5F,0x97,0x44,0x17,0xC4,0xA7,0x7E,0x3D,0x64,0x5D,0x19,0x73},
        {0x60,0x81,0x4F,0xDC,0x22,0x2A,0x90,0x88,0x46,0xEE,0xB8,0x14,0xDE,0x5E,0x0B,0xDB},
        {0xE0,0x32,0x3A,0x0A,0x49,0x06,0x24,0x5C,0xC2,0xD3,0xAC,0x62,0x91,0x95,0xE4,0x79},
        {0xE7,0xC8,0x37,0x6D,0x8D,0xD5,0x4E,0xA9,0x6C,0x56,0xF4,0xEA,0x65,0x7A,0xAE,0x08},
        {0xBA,0x78,0x25,0x2E,0x1C,0xA6,0xB4,0xC6,0xE8,0xDD,0x74,0x1F,0x4B,0xBD,0x8B,0x8A},
        {0x70,0x3E,0xB5,0x66,0x48,0x03,0xF6,0x0E,0x61,0x35,0x57,0xB9,0x86,0xC1,0x1D,0x9E},
        {0xE1,0xF8,0x98,0x11,0x69,0xD9,0x8E,0x94,0x9B,0x1E,0x87,0xE9,0xCE,0x55,0x28,0xDF},
        {0x8C,0xA1,0x89,0x0D,0xBF,0xE6,0x42,0x68,0x41,0x99,0x2D,0x0F,0xB0,0x54,0xBB,0x16}
        };

        /// <summary>
        /// Rcon
        /// </summary>
        private static readonly uint[] Rcon = new uint[]
        {
        0x01000000, 0x02000000,
        0x04000000, 0x08000000,
        0x10000000, 0x20000000,
        0x40000000, 0x80000000,
        0x1B000000, 0x36000000
        };

        /// <summary>
        /// AES算法选择
        /// </summary>
        /// <param name="keySize"></param>
        /// <exception cref="ArgumentException"></exception>
        public AES(int keySize = 128)
        {
            switch (keySize)
            {
                case 128: _nk = 4; _rounds = 10; break;
                case 192: _nk = 6; _rounds = 12; break;
                case 256: _nk = 8; _rounds = 14; break;
                default: throw new ArgumentException("Invalid key size");
            }

        }
        /// <summary>
        /// 加密
        /// </summary>
        /// <param name="input"></param>
        /// <param name="key"></param>
        /// <returns></returns>
        /// <exception cref="ArgumentException"></exception>
        public byte[] Encrypt(byte[] input, byte[] key)
        {
            //对输入的数据进行格式检验
            if (input == null || input.Length != 16)
                throw new ArgumentException("输入的明文长度必须为16字节（128bit）");
            if (key == null || key.Length != _nk * 4)
                throw new ArgumentException("输入的密钥长度必须为{_nk * 4}字节（{_nk} * 32bit）");

            byte[,] state = InitalizeState(input);//初始化状态矩阵
            _roundKeys = KeyExpansion(key);

            AddRoundKey(ref state, 0);// 初始的轮密钥相加

            //主加密轮次
            for (int round = 1; round < _rounds; round++)
            {
                SubBytes(ref state);//字节代换
                ShiftRows(ref state);//行移位
                MixColumns(ref state);//列混合
                AddRoundKey(ref state, round);//轮密钥相加
            }

            //最后一轮加密，没有列混合
            SubBytes(ref state);
            ShiftRows(ref state);
            AddRoundKey(ref state, _rounds);

            return GetStateBytes(state);
        }
        private byte[,] InitalizeState(byte[] input)
        {
            byte[,] state = new byte[4, 4];
            for (int i = 0; i < 4; i++)
                for (int j = 0; j < 4; j++)
                    state[j, i] = input[i * 4 + j];//将输入的字节按列填充
            return state;
        }

        // 密钥扩展
        private uint[] KeyExpansion(byte[] key)
        {
            int totalWords = (_rounds + 1) * 4;// 计算轮密钥的总字数，128为10轮，192为12轮，256为14轮。
            uint[] w = new uint[totalWords];

            //初始密钥填充
            for (int i = 0; i < _nk; i++)
                w[i] = (uint)(key[4 * i] << 24 | key[4 * i + 1] << 16 | key[4 * i + 2] << 8 | key[4 * i + 3]);

            //密钥扩展主循环
            for (int i = _nk; i < totalWords; i++)
            {
                uint temp = w[i - 1];
                /*
                 * 
                 */
                if (i % _nk == 0)
                    temp = SubWord(RotWord(temp)) ^ Rcon[i / _nk - 1];
                else if (_nk > 6 && i % _nk == 4) //AES256的特殊处理
                    temp = SubWord(temp);

                w[i] = w[i - _nk] ^ temp;

            }

            return w;
        }

        private uint RotWord(uint word) => (word >> 24) | (word << 8);//循环左移8位和循环右移24位合并

        private uint SubWord(uint word)
        {
            byte[] bytes = BitConverter.GetBytes(word);
            for (int i = 0; i < 4; i++)
                bytes[i] = Sbox[bytes[i] >> 4, bytes[i] & 0x0F];
            return BitConverter.ToUInt32(bytes, 0);

        }

        private void SubBytes(ref byte[,] state)
        {
            for (int i = 0; i < 4; i++)
                for (int j = 0; j < 4; j++)
                    state[i, j] = Sbox[state[i, j] >> 4, state[i, j] & 0x0F];//高四位行索引，低四位列索引
        }

        //行移位
        private void ShiftRows(ref byte[,] state)
        {
            //第零行不变，第一行左移一位，第二行左移两位，第三行左移三位
            byte temp = state[1, 0];
            state[1, 0] = state[1, 1];
            state[1, 1] = state[1, 2];
            state[1, 2] = state[1, 3];
            state[1, 3] = temp;

            temp = state[2, 0];
            state[2, 0] = state[2, 2];
            state[2, 2] = temp;
            temp = state[2, 1];
            state[2, 1] = state[2, 3];
            state[2, 3] = temp;

            temp = state[3, 0];
            state[3, 0] = state[3, 3];
            state[3, 3] = state[3, 2];
            state[3, 2] = state[3, 1];
            state[3, 1] = temp;
        }

        //列混合
        private void MixColumns(ref byte[,] state)
        {
            for (int col = 0; col < 4; col++)
            {
                byte s0 = state[0, col];
                byte s1 = state[1, col];
                byte s2 = state[2, col];
                byte s3 = state[3, col];

                state[0, col] = (byte)(Gmul(2, s0) ^ Gmul(3, s1) ^ s2 ^ s3);
                state[1, col] = (byte)(s0 ^ Gmul(2, s1) ^ Gmul(3, s2) ^ s3);
                state[2, col] = (byte)(s0 ^ s1 ^ Gmul(2, s2) ^ Gmul(3, s3));
                state[3, col] = (byte)(Gmul(3, s0) ^ s1 ^ s2 ^ Gmul(2, s3));
            }
        }

        //GF(2^8)乘法
        private static byte Gmul(byte a, byte b)
        {
            byte p = 0;
            for (int i = 0; i < 8; i++)
            {
                if ((b & 1) != 0) p ^= a;
                bool hi_bit_set = (a & 0x80) != 0;
                a <<= 1;
                if (hi_bit_set) a ^= 0x1b;
                b >>= 1;
            }
            return p;
        }

        //轮密钥相加
        private void AddRoundKey(ref byte[,] state, int round)
        {
            for (int col = 0; col < 4; col++)
            {
                uint word = _roundKeys[round * 4 + col];
                state[0, col] ^= (byte)(word >> 24);
                state[1, col] ^= (byte)(word >> 16);
                state[2, col] ^= (byte)(word >> 8);
                state[3, col] ^= (byte)word;
            }
        }

        //将状态矩阵按列填充到输出字节数组
        private byte[] GetStateBytes(byte[,] state)
        {
            byte[] output = new byte[16];
            for (int i = 0; i < 4; i++)
                for (int j = 0; j < 4; j++)
                    output[i * 4 + j] = state[j, i];
            return output;
        }

        //增加一个返回加密一万轮每轮第一个字节功能
        public byte GetFirstSubByte(byte[] input, byte[] key)
        {
            var state = InitalizeState(input);
            _roundKeys = KeyExpansion(key);

            // 只执行第一轮的SubBytes
            SubBytes(ref state);

            // 返回第一个字节
            return state[0, 0];
        }


        //增加一个解密功能
        public byte[] Decrypt(byte[] input, byte[] key)
        {
            // 参数验证
            if (input == null || input.Length != 16)
                throw new ArgumentException("Input must be 16 bytes (128 bits)");
            if (key == null || key.Length != _nk * 4)
                throw new ArgumentException($"Key must be {_nk * 4} bytes ({_nk * 32} bits)");

            byte[,] state = InitalizeState(input);
            _roundKeys = KeyExpansion(key);

            // 初始轮密钥加
            AddRoundKey(ref state, _rounds);

            // 主解密轮次
            for (int round = _rounds - 1; round >= 1; round--)
            {
                InvShiftRows(ref state);
                InvSubBytes(ref state);
                AddRoundKey(ref state, round);
                InvMixColumns(ref state);
            }

            // 最终轮（无InvMixColumns）
            InvShiftRows(ref state);
            InvSubBytes(ref state);
            AddRoundKey(ref state, 0);

            return GetStateBytes(state);
        }
        //逆s盒
        private static readonly byte[,] InvSBox = new byte[16, 16]
        {
    {0x52,0x09,0x6A,0xD5,0x30,0x36,0xA5,0x38,0xBF,0x40,0xA3,0x9E,0x81,0xF3,0xD7,0xFB},
    {0x7C,0xE3,0x39,0x82,0x9B,0x2F,0xFF,0x87,0x34,0x8E,0x43,0x44,0xC4,0xDE,0xE9,0xCB},
    {0x54,0x7B,0x94,0x32,0xA6,0xC2,0x23,0x3D,0xEE,0x4C,0x95,0x0B,0x42,0xFA,0xC3,0x4E},
    {0x08,0x2E,0xA1,0x66,0x28,0xD9,0x24,0xB2,0x76,0x5B,0xA2,0x49,0x6D,0x8B,0xD1,0x25},
    {0x72,0xF8,0xF6,0x64,0x86,0x68,0x98,0x16,0xD4,0xA4,0x5C,0xCC,0x5D,0x65,0xB6,0x92},
    {0x6C,0x70,0x48,0x50,0xFD,0xED,0xB9,0xDA,0x5E,0x15,0x46,0x57,0xA7,0x8D,0x9D,0x84},
    {0x90,0xD8,0xAB,0x00,0x8C,0xBC,0xD3,0x0A,0xF7,0xE4,0x58,0x05,0xB8,0xB3,0x45,0x06},
    {0xD0,0x2C,0x1E,0x8F,0xCA,0x3F,0x0F,0x02,0xC1,0xAF,0xBD,0x03,0x01,0x13,0x8A,0x6B},
    {0x3A,0x91,0x11,0x41,0x4F,0x67,0xDC,0xEA,0x97,0xF2,0xCF,0xCE,0xF0,0xB4,0xE6,0x73},
    {0x96,0xAC,0x74,0x22,0xE7,0xAD,0x35,0x85,0xE2,0xF9,0x37,0xE8,0x1C,0x75,0xDF,0x6E},
    {0x47,0xF1,0x1A,0x71,0x1D,0x29,0xC5,0x89,0x6F,0xB7,0x62,0x0E,0xAA,0x18,0xBE,0x1B},
    {0xFC,0x56,0x3E,0x4B,0xC6,0xD2,0x79,0x20,0x9A,0xDB,0xC0,0xFE,0x78,0xCD,0x5A,0xF4},
    {0x1F,0xDD,0xA8,0x33,0x88,0x07,0xC7,0x31,0xB1,0x12,0x10,0x59,0x27,0x80,0xEC,0x5F},
    {0x60,0x51,0x7F,0xA9,0x19,0xB5,0x4A,0x0D,0x2D,0xE5,0x7A,0x9F,0x93,0xC9,0x9C,0xEF},
    {0xA0,0xE0,0x3B,0x4D,0xAE,0x2A,0xF5,0xB0,0xC8,0xEB,0xBB,0x3C,0x83,0x53,0x99,0x61},
    {0x17,0x2B,0x04,0x7E,0xBA,0x77,0xD6,0x26,0xE1,0x69,0x14,0x63,0x55,0x21,0x0C,0x7D}

};

        // 逆行移位
        private void InvShiftRows(ref byte[,] state)
        {
            // 第1行右移1字节
            byte temp = state[1, 3];
            state[1, 3] = state[1, 2];
            state[1, 2] = state[1, 1];
            state[1, 1] = state[1, 0];
            state[1, 0] = temp;

            // 第2行右移2字节
            temp = state[2, 2];
            state[2, 2] = state[2, 0];
            state[2, 0] = temp;
            temp = state[2, 3];
            state[2, 3] = state[2, 1];
            state[2, 1] = temp;

            // 第3行右移3字节（等价左移1）
            temp = state[3, 0];
            state[3, 0] = state[3, 1];
            state[3, 1] = state[3, 2];
            state[3, 2] = state[3, 3];
            state[3, 3] = temp;
        }

        // 逆字节替换
        private void InvSubBytes(ref byte[,] state)
        {
            for (int i = 0; i < 4; i++)
                for (int j = 0; j < 4; j++)
                    state[i, j] = InvSBox[state[i, j] >> 4, state[i, j] & 0x0F];
        }

        // 逆列混合
        private void InvMixColumns(ref byte[,] state)
        {
            for (int col = 0; col < 4; col++)
            {
                byte s0 = state[0, col];
                byte s1 = state[1, col];
                byte s2 = state[2, col];
                byte s3 = state[3, col];

                state[0, col] = (byte)(Gmul(0x0E, s0) ^ Gmul(0x0B, s1) ^ Gmul(0x0D, s2) ^ Gmul(0x09, s3));
                state[1, col] = (byte)(Gmul(0x09, s0) ^ Gmul(0x0E, s1) ^ Gmul(0x0B, s2) ^ Gmul(0x0D, s3));
                state[2, col] = (byte)(Gmul(0x0D, s0) ^ Gmul(0x09, s1) ^ Gmul(0x0E, s2) ^ Gmul(0x0B, s3));
                state[3, col] = (byte)(Gmul(0x0B, s0) ^ Gmul(0x0D, s1) ^ Gmul(0x09, s2) ^ Gmul(0x0E, s3));
            }
        }


    }
}
